import os
import json
import random
from collections import defaultdict
from typing import List, Dict, Optional

import datasets
from langchain.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field
import concurrent.futures
from tqdm import tqdm
import time


# 定义结构化输出模型
class QuestionAnswer(BaseModel):
    answer: str = Field(description="模拟真实的客户发言文本")


# 提示模板定义
prompt_template = PromptTemplate(
    input_variables=["question"],
    template="Answer the following question:\n\nQuestion: {question}\n\nAnswer:"
)


def get_structured_output(model_class, prompt: PromptTemplate, question: str):
    """
    使用 LLM 获取结构化输出。
    """
    os.environ["OPENAI_API_KEY"] = 'local-qwen2.5-72b-little-brother'
    base_url = "http://127.0.0.1:8814/v1"
    model_name = "qwq-32b-preview"

    # 添加重试机制
    max_retries = 3
    for attempt in range(max_retries):
        try:
            llm = ChatOpenAI(
                base_url=base_url,
                model=model_name,
                api_key=os.environ["OPENAI_API_KEY"],
                max_tokens=1000,
                temperature=0.7
            )

            chain = prompt | llm.with_structured_output(model_class)
            return chain.invoke({"question": question})
        except Exception as e:
            if attempt < max_retries - 1:
                time.sleep(2 ** attempt)  # 指数退避
                continue
            else:
                raise e


def load_intention_data(csv_path: str) -> Dict[str, List[str]]:
    """
    加载意图识别数据，按类别分组。

    Args:
        csv_path (str): CSV 文件路径。

    Returns:
        Dict[str, List[str]]: 按意图类别分组的提示列表。
    """
    dataset = datasets.load_dataset('csv', data_files=csv_path)['train']
    intentions = defaultdict(list)
    
    for item in dataset:
        intentions[item['answer']].append(item['prompt'])
    
    return dict(intentions)


def build_prompt_for_intention(few_shots: List[str], intention_type: str) -> str:
    """
    构建用于生成特定意图类型的客户话语的提示语。

    Args:
        few_shots (List[str]): few-shot 示例列表。
        intention_type (str): 意图类型。

    Returns:
        str: 构建好的提示字符串。
    """
    # 组装 few-shot 示例为带案例标题的格式
    examples_str = ""
    for i, ex in enumerate(few_shots, 3):
        examples_str += f"## 案例 {i}\n{ex}\n\n"

    intention_descriptions = {
        "PRODUCT_QA": "关于保险产品具体问题的询问，如产品保障范围、条款细节、投保规则等",
        "PRODUCT_ALL_CONTENT": "询问保险产品的完整保障内容，如保障责任、除外责任等",
        "PRODUCT_REC": "请求推荐合适的保险产品，如根据年龄、健康状况等条件推荐产品",
        "OTHER": "其他类型的问题或陈述，如基本咨询、一般性问题等",
        "APPOINTMENT_CONSULTATION": "预约人工咨询服务，如要求转接人工客服",
        "PRODUCT_HB": "关于保险产品条款的询问，如保险责任、免责条款等",
        "PRODUCT_HEALTH": "关于健康保险的询问，如健康告知、核保规则等",
        "PRODUCT_HB_REC": "请求推荐意外险产品",
        "PRODUCT_FEE": "关于保险费用的询问，如保费计算、费率表等",
        "PRODUCT_BUY": "关于购买保险流程的询问，如投保链接、购买方式等",
        "PRODUCT_COMPARE": "比较不同保险产品，如优缺点对比、性价比分析等",
        "PREMIUM_CALCULATION": "保险费用计算相关问题，如具体保费金额查询",
        "INSURANCE_COMPANY": "关于保险公司的问题，如公司背景、偿付能力等",
        "PRODUCT_PLAN": "关于保险计划的问题，如家庭保险规划、个人保障方案等",
        "INSURANCE_KNOWLEDGE": "保险知识相关问题，如保险术语解释、保险原理等",
        "PRODUCT_DUTY": "关于保险责任的问题，如具体保障范围、赔付条件等",
        "BENEFIT_ILLUSTRATION": "关于保险利益说明的问题，如现金价值、年金领取等"
    }
    
    intention_desc = intention_descriptions.get(intention_type, "特定类型的保险相关问题")

    return f"""
# 任务目标：
生成客户关于{intention_desc}的对话内容，只包含客户首句话。

# 要求：
- 生成的语言应自然、口语化，符合真实客户沟通场景；
- 内容必须与{intention_desc}高度相关；
- 可以包含疑问句、陈述句、情绪化表达、方言等；
- 不包括骂人、投诉等负面话语；
- 问题应该具体明确，避免过于宽泛的提问；
- 每个生成的案例应该具有独特性，避免重复模式；

# 输出格式：
- 返回 JSON 格式，字段包括 [answer]
- [answer]字段中只包含客户说的多句话，每条用换行分隔

# Few-Shot 示例：
{examples_str.strip()}
## 新案例：
请生成一个符合上述要求的客户咨询案例。
"""


def generate_intention_data(dialog: str, few_shots: List[str], intention_type: str) -> Dict:
    """
    生成特定意图类型的客户话语。

    Args:
        dialog (str): 输入内容（仅占位）
        few_shots (List[str]): few-shot 示例
        intention_type (str): 意图类型

    Returns:
        dict: 包含生成结果的字典。
    """
    try:
        prompt = build_prompt_for_intention(few_shots, intention_type)
        response = get_structured_output(QuestionAnswer, prompt_template, prompt)
        return response.model_dump()
    except Exception as e:
        print(f"Error generating for intention '{intention_type}': {e}")
        return {"intention": intention_type, "error": str(e)}


def process_intention_item(intention_type: str, few_shots: List[str]) -> Dict:
    """
    处理单个意图类型的数据生成。

    Args:
        intention_type (str): 意图类型
        few_shots (List[str]): few-shot 示例

    Returns:
        dict: 包含生成结果的字典。
    """
    return generate_intention_data("", few_shots, intention_type)


def generate_intention_samples(
        intention_data: Dict[str, List[str]],
        samples_per_category: int = 100,
        save_dir: str = "/Users/sd/Desktop/mycode/myalgo/milvus/data-gen/zero_gen"
) -> Dict[str, List[Dict]]:
    """
    为每个意图类别生成扩充数据。

    Args:
        intention_data (Dict[str, List[str]]): 按意图类别分组的原始数据
        samples_per_category (int): 每个类别生成的样本数量
        save_dir (str): 保存目录

    Returns:
        Dict[str, List[Dict]]: 所有生成的意图数据
    """
    all_results = {}
    
    for intention_type, examples in intention_data.items():
        print(f"Generating data for intention: {intention_type} ({len(examples)} examples)")
        
        results = []
        save_path = os.path.join(save_dir, f"{intention_type}_generated.json")
        
        # 如果该类别样本数少于3个，则跳过
        if len(examples) < 3:
            print(f"Skipping {intention_type} - too few examples ({len(examples)})")
            continue
            
        # 减少并发数以降低连接错误概率
        with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
            # 为每个任务选择随机的few-shot示例
            futures = {}
            for i in range(samples_per_category):
                # 随机选择few-shot示例，最多3个
                sample_size = min(3, len(examples))
                few_shots = random.sample(examples, sample_size)
                future = executor.submit(process_intention_item, intention_type, few_shots)
                futures[future] = i

            for future in tqdm(concurrent.futures.as_completed(futures), 
                             total=len(futures), 
                             desc=f"Generating {intention_type}", 
                             colour="GREEN"):
                try:
                    result = future.result(timeout=60)
                    results.append(result)
                except Exception as e:
                    print(f"Task failed for {intention_type} with exception: {e}")

        # 保存单个类别的结果
        with open(save_path, 'w', encoding='utf-8') as f:
            json.dump(results, f, indent=2, ensure_ascii=False)
            
        all_results[intention_type] = results
        print(f"Generated {len(results)} samples for {intention_type}")
        
    return all_results


if __name__ == "__main__":
    # 设置参数
    csv_path = "/Users/sd/Desktop/mycode/myalgo/milvus/data-gen/zero_gen/intention.csv"
    save_dir = "/Users/sd/Desktop/mycode/myalgo/milvus/data-gen/zero_gen"
    
    # 确保保存目录存在
    os.makedirs(save_dir, exist_ok=True)
    
    # 加载意图数据
    intention_data = load_intention_data(csv_path)
    print(f"Loaded {len(intention_data)} intention categories")
    for intention_type, examples in intention_data.items():
        print(f"- {intention_type}: {len(examples)} examples")
    
    # 为每个意图类别生成数据
    generated_data = generate_intention_samples(
        intention_data=intention_data,
        samples_per_category=50,  # 每个类别生成50个样本以减少运行时间
        save_dir=save_dir
    )
    
    # 保存汇总结果
    summary_path = os.path.join(save_dir, "all_intentions_generated.json")
    with open(summary_path, 'w', encoding='utf-8') as f:
        json.dump(generated_data, f, indent=2, ensure_ascii=False)
        
    print("所有意图数据生成完成！")